#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
生理数据可视化程序
根据心电图和皮电数据生成时间序列图表
包含平均心率、心率变异性(LF/HF)、SCL三个指标的可视化
"""

import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import numpy as np
import os
from matplotlib import rcParams

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei']
plt.rcParams['axes.unicode_minus'] = False

class PhysioDataVisualizer:
    def __init__(self, base_path):
        """
        初始化可视化器
        
        Args:
            base_path: 基础路径
        """
        self.base_path = base_path
        # 修改路径指向可视化结果目录
        self.group_data_path = os.path.join(base_path, "群体平均值数据.csv")
        self.individual_data_path = os.path.join(base_path, "以人名为依据的数据.csv")
        
        # 创建输出目录
        self.output_dir = os.path.join(base_path, "可视化结果", "图片")
        
        self.load_data()
        
    def load_data(self):
        """加载数据文件"""
        try:
            # 加载群体平均值数据
            self.group_data = pd.read_csv(self.group_data_path, encoding='utf-8')
            print(f"群体数据加载成功: {len(self.group_data)} 条记录")
            
            # 加载个体数据
            self.individual_data = pd.read_csv(self.individual_data_path, encoding='utf-8')
            print(f"个体数据加载成功: {len(self.individual_data)} 条记录")
            
            # 数据清洗和预处理
            self.preprocess_data()
            
        except Exception as e:
            print(f"数据加载失败: {e}")
            raise
    
    def preprocess_data(self):
        """数据预处理"""
        # 跳过标题行（如果存在）
        if self.group_data.iloc[0, 0] == 'stage_id':
            self.group_data = self.group_data.iloc[1:].reset_index(drop=True)
        
        if self.individual_data.iloc[0, 0] == 'stage_id':
            self.individual_data = self.individual_data.iloc[1:].reset_index(drop=True)
        
        # 转换数据类型
        self.group_data = self.group_data.apply(pd.to_numeric, errors='ignore')
        self.individual_data = self.individual_data.apply(pd.to_numeric, errors='ignore')
        
        # 检查关键列是否存在
        required_group_cols = ['平均心率均值(bpm)', 'LF/HF比值均值', 'SCL均值(μS)_mean']
        missing_cols = [col for col in required_group_cols if col not in self.group_data.columns]
        if missing_cols:
            print(f"群体数据缺少列: {missing_cols}")
            print(f"可用列: {list(self.group_data.columns)}")
        
        required_individual_cols = ['平均心率(bpm)', 'LF/HF比值', 'SCL均值(μS)']
        missing_individual_cols = [col for col in required_individual_cols if col not in self.individual_data.columns]
        if missing_individual_cols:
            print(f"个体数据缺少列: {missing_individual_cols}")
            print(f"可用列: {list(self.individual_data.columns)}")
        
        # 处理缺失值（使用正确的列名）
        if not missing_cols:
            self.group_data = self.group_data.dropna(subset=required_group_cols)
        
        print("数据预处理完成")
    
    def calculate_individual_statistics(self, stage_id, group_type):
        """
        计算个体数据的统计量，用于生成阴影区域
        
        Args:
            stage_id: 阶段ID
            group_type: 分组类型
            
        Returns:
            dict: 包含统计量的字典
        """
        # 过滤数据
        filtered_data = self.individual_data[
            (self.individual_data['阶段ID'] == stage_id) & 
            (self.individual_data['分组类型'] == group_type)
        ]
        
        if filtered_data.empty:
            return None
        
        # 按分段索引分组计算统计量
        stats = {}
        for segment_idx in sorted(filtered_data['分段索引'].unique()):
            segment_data = filtered_data[filtered_data['分段索引'] == segment_idx]
            
            if len(segment_data) > 0:
                stats[segment_idx] = {
                    'mhr_mean': segment_data['平均心率(bpm)'].mean(),
                    'mhr_std': segment_data['平均心率(bpm)'].std(),
                    'lf_hf_mean': segment_data['LF/HF比值'].mean(),
                    'lf_hf_std': segment_data['LF/HF比值'].std(),
                    'scl_mean': segment_data['SCL均值(μS)'].mean(),
                    'scl_std': segment_data['SCL均值(μS)'].std(),
                    'count': len(segment_data)
                }
        
        return stats
    
    def create_time_series_plot(self, stage_id, group_type, metric, save_path=None, export_csv=True):
        """
        创建时间序列图表
        
        Args:
            stage_id: 阶段ID
            group_type: 分组类型
            metric: 指标类型 ('heart_rate', 'lf_hf', 'scl')
            save_path: 保存路径（可选）
            export_csv: 是否导出底层数据为CSV（默认True）
        """
        # 过滤群体数据
        group_filtered = self.group_data[
            (self.group_data['阶段ID'] == stage_id) & 
            (self.group_data['分组类型'] == group_type)
        ].sort_values('分段索引')
        
        if group_filtered.empty:
            print(f"未找到阶段{stage_id}，{group_type}的数据")
            return
        
        # 获取个体数据统计量
        individual_stats = self.calculate_individual_statistics(stage_id, group_type)
        
        # 设置指标相关参数
        metric_configs = {
            'heart_rate': {
                'group_mean_col': '平均心率均值(bpm)',
                'group_std_col': '平均心率标准差(bpm)',
                'individual_mean_key': 'mhr_mean',
                'individual_std_key': 'mhr_std',
                'title': f'平均心率 - 阶段{stage_id} - {group_type}',
                'ylabel': '平均心率 (bpm)',
                'color': '#66C2A5'  # 用户指定颜色
            },
            'lf_hf': {
                'group_mean_col': 'LF/HF比值均值',
                'group_std_col': 'LF/HF比值标准差',
                'individual_mean_key': 'lf_hf_mean',
                'individual_std_key': 'lf_hf_std',
                'title': f'心率变异性(LF/HF比值) - 阶段{stage_id} - {group_type}',
                'ylabel': 'LF/HF比值',
                'color': '#66C2A5'  # 调整为用户指定的颜色
            },
            'scl': {
                'group_mean_col': 'SCL均值(μS)_mean',
                'group_std_col': 'SCL均值(μS)_std',
                'individual_mean_key': 'scl_mean',
                'individual_std_key': 'scl_std',
                'title': f'皮肤电导水平(SCL) - 阶段{stage_id} - {group_type}',
                'ylabel': 'SCL (μS)',
                'color': '#66C2A5'  # 调整为用户指定的颜色

            }
        }
        
        config = metric_configs[metric]
        
        # 设置图表为瘦长布局
        fig, ax = plt.subplots(figsize=(10, 5))  # 调整为瘦长比例
        
        # 准备时间轴数据（每30秒一个数据点）
        segments = group_filtered['分段索引'].values
        time_points = (segments - 1) * 30  # 起点对齐：索引1→0秒
        group_means = group_filtered[config['group_mean_col']].values
        
        # 过滤掉NaN值的数据点
        valid_group_mask = ~np.isnan(group_means)
        time_points_valid = time_points[valid_group_mask]
        group_means_valid = group_means[valid_group_mask]
        
        # **动态设置X轴最大值 - 方案A：不留白，基于实际数据范围**
        if len(time_points_valid) > 0:
            max_time = max(time_points_valid)  # 直接使用最后一个数据点，不加边距
            # 确保X轴最小为30秒（1个时间段）
            max_time = max(max_time, 30)
        else:
            max_time = 30  # 没有数据时默认30秒
        
        # **修复问题1：群体平均值（中线）- 只绘制有效数据点，确保有线条显示**
        if len(time_points_valid) > 0:
            ax.plot(time_points_valid, group_means_valid, '-', color=config['color'], linewidth=2, label='群体平均值', marker='o', markersize=4)
        else:
            # 如果没有有效的群体数据，显示提示信息
            ax.text(max_time/2, 0.5, '无有效群体数据', ha='center', va='center', fontsize=12, color='red')
        
        # 收集所有有效的Y值用于设置Y轴范围
        all_y_values = []
        if len(group_means_valid) > 0:
            all_y_values.extend(group_means_valid)
        
        # 如果有个体数据，添加个体数据的阴影
        if individual_stats:
            individual_segments = sorted(individual_stats.keys())
            individual_time_points = (np.array(individual_segments) - 1) * 30
            individual_means = [individual_stats[seg][config['individual_mean_key']] for seg in individual_segments]
            individual_stds = [individual_stats[seg][config['individual_std_key']] for seg in individual_segments]
            
            # 过滤NaN值
            individual_time_points = np.array(individual_time_points)
            individual_means = np.array(individual_means)
            individual_stds = np.array(individual_stds)
            
            # **修复问题1和3：更严格的NaN值处理**
            valid_individual_mask = ~(np.isnan(individual_means) | np.isnan(individual_stds) | (individual_stds == 0))
            
            if valid_individual_mask.any():
                valid_time_points = individual_time_points[valid_individual_mask]
                valid_means = individual_means[valid_individual_mask]
                valid_stds = individual_stds[valid_individual_mask]
                
                # 个体数据平均值（虚线）
                ax.plot(valid_time_points, valid_means, '--', color=config['color'], 
                       linewidth=1.5, label='个体数据平均值', marker='s', markersize=3)
                
                # **修复问题3&4：个体数据标准差阴影 - 计算上下边界，确保不越界**
                upper_bound_ind = valid_means + valid_stds
                lower_bound_ind = valid_means - valid_stds
                
                # 添加Y值用于范围计算
                all_y_values.extend(upper_bound_ind)
                all_y_values.extend(lower_bound_ind)
                
                ax.fill_between(valid_time_points, lower_bound_ind, upper_bound_ind, 
                              alpha=0.2, color="#FC8D62", label='个体标准差范围')
        
        # 设置图表属性
        ax.set_xlabel('时间 (秒)', fontsize=11)
        ax.set_ylabel(config['ylabel'], fontsize=11)
        ax.set_title(config['title'], fontsize=13, fontweight='bold', pad=20)
        ax.grid(True, alpha=0.2, linestyle='--')
        ax.legend(fontsize=9, loc='upper right')
        
        # **设置X轴范围 - 动态且不留白（末端与最后一个点对齐）**
        ax.set_xlim(0, max_time)
        ax.set_xticks(np.arange(0, max_time + 1, 30))  # 每30秒一个刻度，末端对齐
        
        # **修复问题3&4：设置Y轴范围 - 适度扩大显示范围，让阴影面积在视觉上更合理**
        if all_y_values:
            y_min_data = min(all_y_values)
            y_max_data = max(all_y_values)
            
            # 根据指标类型设置合适的显示范围，避免差距过大
            if metric == 'heart_rate':
                # 心率：固定在45–120区间内，适度边距
                y_range = y_max_data - y_min_data
                margin = max(0.18 * y_range, 6.0)
                y_min = max(45, y_min_data - margin)
                y_max = min(120, y_max_data + margin)
            elif metric == 'lf_hf':
                # LF/HF：上限固定到20，适度边距
                y_range = y_max_data - y_min_data
                margin = max(0.22 * y_range, 1.2)
                y_min = max(0, y_min_data - margin)
                y_max = min(20, y_max_data + margin)
            elif metric == 'scl':
                # SCL：上限固定到35，适度边距
                y_range = y_max_data - y_min_data
                margin = max(0.22 * y_range, 1.8)
                y_min = max(0, y_min_data - margin)
                y_max = min(35, y_max_data + margin)
            else:
                # 其他指标：适度增加边距
                y_range = y_max_data - y_min_data
                margin = max(0.2 * y_range, 1.0)
                y_min = y_min_data - margin
                y_max = y_max_data + margin
            
            ax.set_ylim(y_min, y_max)
        else:
            # 没有有效数据时的默认范围（方案A对应）
            if metric == 'heart_rate':
                ax.set_ylim(45, 120)
            elif metric == 'lf_hf':
                ax.set_ylim(0, 20)
            elif metric == 'scl':
                ax.set_ylim(0, 35)
            else:
                ax.set_ylim(0, 1)
        
        # 调整布局，减少空白
        plt.tight_layout()
        plt.subplots_adjust(top=0.9, bottom=0.15, left=0.08, right=0.95)
        
        # 导出底层数据为CSV
        if export_csv and save_path:
            try:
                # 合并群体与个体的所有分段索引
                segments_union = sorted(set(segments.tolist()) | (set(individual_stats.keys()) if individual_stats else set()))
                rows = []
                for seg in segments_union:
                    time_sec = (seg - 1) * 30
                    grp_row = group_filtered[group_filtered['分段索引'] == seg]
                    grp_mean = grp_row[config['group_mean_col']].values[0] if len(grp_row) > 0 else np.nan
                    ind_mean = individual_stats.get(seg, {}).get(config['individual_mean_key'], np.nan) if individual_stats else np.nan
                    ind_std = individual_stats.get(seg, {}).get(config['individual_std_key'], np.nan) if individual_stats else np.nan
                    cnt = individual_stats.get(seg, {}).get('count', 0) if individual_stats else 0
                    lower = ind_mean - ind_std if (not np.isnan(ind_mean) and not np.isnan(ind_std)) else np.nan
                    upper = ind_mean + ind_std if (not np.isnan(ind_mean) and not np.isnan(ind_std)) else np.nan
                    
                    rows.append({
                        'stage_id': stage_id,
                        'group_type': group_type,
                        'metric': metric,
                        'segment_index': seg,
                        'time_sec': time_sec,
                        'group_mean': grp_mean,
                        'individual_mean': ind_mean,
                        'individual_std': ind_std,
                        'lower': lower,
                        'upper': upper,
                        'count': cnt
                    })
                
                csv_data = pd.DataFrame(rows)
                csv_path = save_path.replace('.png', '.csv')
                csv_data.to_csv(csv_path, index=False, encoding='utf-8-sig')
                
            except Exception as e:
                print(f"导出CSV时出错: {e}")
        
        # 保存图片
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight', facecolor='white')
            print(f"图表已保存: {save_path}")
        
        return fig, ax
    
    def generate_all_plots(self, output_dir="visualization_output"):
        """
        生成所有的可视化图表
        
        Args:
            output_dir: 输出目录
        """
        # 创建输出目录
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        
        # 获取所有的阶段和分组组合
        stage_groups = self.group_data[['阶段ID', '分组类型']].drop_duplicates()
        
        metrics = ['heart_rate', 'lf_hf', 'scl']
        metric_names = ['平均心率', '心率变异性', 'SCL']
        
        total_plots = len(stage_groups) * len(metrics)
        current_plot = 0
        
        print(f"开始生成 {total_plots} 个图表...")
        
        for _, row in stage_groups.iterrows():
            stage_id = row['阶段ID']
            group_type = row['分组类型']
            
            print(f"\n处理阶段{stage_id} - {group_type}:")
            
            for metric, metric_name in zip(metrics, metric_names):
                current_plot += 1
                
                try:
                    # 生成文件名
                    safe_group_type = group_type.replace('/', '_')
                    filename = f"阶段{stage_id}_{safe_group_type}_{metric_name}.png"
                    save_path = os.path.join(output_dir, filename)
                    
                    # 创建图表并导出CSV
                    fig, ax = self.create_time_series_plot(stage_id, group_type, metric, save_path, export_csv=True)
                    
                    # 关闭图表释放内存
                    plt.close(fig)
                    
                    print(f"  [{current_plot}/{total_plots}] {metric_name} 图表及CSV已生成")
                    
                except Exception as e:
                    print(f"  生成 {metric_name} 图表时出错: {e}")
        
        print(f"\n所有图表及CSV生成完成！保存在目录: {output_dir}")
    
    def display_data_summary(self):
        """显示数据摘要信息"""
        print("\n=== 数据摘要 ===")
        print(f"群体数据: {len(self.group_data)} 条记录")
        print(f"个体数据: {len(self.individual_data)} 条记录")
        
        # 显示阶段和分组信息
        stage_groups = self.group_data[['阶段ID', '分组类型']].drop_duplicates()
        print("\n可用的阶段和分组组合:")
        for _, row in stage_groups.iterrows():
            stage_id = row['阶段ID']
            group_type = row['分组类型']
            count = len(self.group_data[
                (self.group_data['阶段ID'] == stage_id) & 
                (self.group_data['分组类型'] == group_type)
            ])
            print(f"  阶段{stage_id} - {group_type}: {count} 个时间段")


def main():
    """主函数"""
    # 设置数据路径
    base_path = r"d:\python代码\心理所\冯亚楠\心电图预处理\预处理文件"
    
    try:
        # 创建可视化器
        visualizer = PhysioDataVisualizer(base_path)
        
        # 显示数据摘要
        visualizer.display_data_summary()
        
        # 生成所有图表
        output_dir = os.path.join(base_path, "可视化结果")
        visualizer.generate_all_plots(output_dir)
        
        print("\n可视化完成！")
        
    except Exception as e:
        print(f"程序执行出错: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()